import os
from typing import Optional, List
import httpx
from mcp.server.fastmcp import FastMCP


def get_api_key() -> str:
    """Get API key from environment variables"""
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("DASHSCOPE_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


mcp = FastMCP("Ali Media MCP Server")


@mcp.tool()
async def create_image(prompt: str, model: str = "wan2.2-t2i-plus") -> str:
    """
    使用阿里云DashScope API创建AI图像生成任务

    参数:
        prompt: 描述生成图像内容的文本提示
        model: 模型名称。示例值：wan2.2-t2i-turbo、wan2.2-t2i-plus。

    返回:
        任务ID，用于后续查询生成结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    CREATE_API = (
        "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
    )

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": model,
        "input": {"prompt": prompt},
        "parameters": {"size": "1024*1024", "n": 1},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(CREATE_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def local_style_transfer(prompt: str, base_image_url: str) -> str:
    """
    使用阿里云DashScope API实现图像局部风格迁移

    参数:
        prompt: 风格描述文本(如："把房子变成木板风格")
        base_image_url: 原始图片URL

    返回:
        任务ID，用于查询处理结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    STYLE_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "stylization_local",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(STYLE_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def description_edit(
    prompt: str, base_image_url: str, strength: float = 0.5
) -> str:
    """
    使用阿里云DashScope API根据文本描述编辑图像

    参数:
        prompt: 编辑要求描述(如："将女孩穿的衣服修改为彩色印花的沙滩衬衣")
        base_image_url: 原始图片URL
        strength: 编辑强度(0-1)，默认0.5

    返回:
        任务ID，用于查询编辑结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    EDIT_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "description_edit",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1, "strength": strength},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(EDIT_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def remove_watermark(prompt: str, base_image_url: str) -> str:
    """
    使用阿里云DashScope API去除图像中的水印/文字

    参数:
        prompt: 去除内容描述(如："去除图像中的文字")
        base_image_url: 待处理图片URL

    返回:
        任务ID，用于查询处理结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    WATERMARK_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "remove_watermark",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(WATERMARK_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def expand_image(
    prompt: str,
    base_image_url: str,
    top_scale: float = 1.0,
    bottom_scale: float = 1.0,
    left_scale: float = 1.0,
    right_scale: float = 1.0,
) -> str:
    """
    使用阿里云DashScope API扩展图像边界

    参数:
        prompt: 扩展内容描述(如："一家人在公园草坪上")
        base_image_url: 原始图片URL
        top_scale: 顶部扩展比例(默认1.0)
        bottom_scale: 底部扩展比例(默认1.0)
        left_scale: 左侧扩展比例(默认1.0)
        right_scale: 右侧扩展比例(默认1.0)

    返回:
        任务ID，用于查询扩展结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    EXPAND_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "expand",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {
            "top_scale": top_scale,
            "bottom_scale": bottom_scale,
            "left_scale": left_scale,
            "right_scale": right_scale,
            "n": 1,
        },
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(EXPAND_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def super_resolution(
    prompt: str, base_image_url: str, upscale_factor: int = 1
) -> str:
    """
    使用阿里云DashScope API增强图像分辨率

    参数:
        prompt: 增强效果描述(如："图像超分")
        base_image_url: 原始图片URL
        upscale_factor: 缩放因子(1-4)，默认1(仅增强)

    返回:
        任务ID，用于查询处理结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    SUPER_RES_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "super_resolution",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {
            "upscale_factor": upscale_factor,
            "n": 1,
        },
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(SUPER_RES_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def colorize_image(prompt: str, base_image_url: str) -> str:
    """
    使用阿里云DashScope API为灰度图像上色

    参数:
        prompt: 上色效果描述(如："蓝色背景，黄色的叶子")
        base_image_url: 原始灰度图片URL

    返回:
        任务ID，用于查询处理结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    COLOR_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "colorization",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(COLOR_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def doodle_to_image(
    prompt: str, base_image_url: str, is_sketch: bool = False
) -> str:
    """
    使用阿里云DashScope API从草图生成图像

    参数:
        prompt: 生成图像描述(如："一颗树，二次元动漫风格")
        base_image_url: 原始草图URL
        is_sketch: 是否为草图输入(默认False)

    返回:
        任务ID，用于查询生成结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    DOODLE_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "doodle",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1, "is_sketch": is_sketch},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(DOODLE_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def cartoon_reference(prompt: str, base_image_url: str) -> str:
    """
    使用阿里云DashScope API基于参考图生成卡通图像

    参数:
        prompt: 卡通动作描述(如："卡通形象小心翼翼地探出头")
        base_image_url: 参考卡通图片URL

    返回:
        任务ID，用于查询生成结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    CARTOON_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "control_cartoon_feature",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(CARTOON_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def global_style_transfer(
    prompt: str, base_image_url: str, strength: float = 0.5
) -> str:
    """
    使用阿里云DashScope API应用全局风格迁移

    参数:
        prompt: 风格描述(如："转换成法国绘本风格")
        base_image_url: 原始图片URL
        strength: 风格迁移强度(0.0-1.0)，默认0.5

    返回:
        任务ID，用于查询处理结果

    异常:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    GLOBAL_STYLE_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-imageedit",
        "input": {
            "function": "stylization_all",
            "prompt": prompt,
            "base_image_url": base_image_url,
        },
        "parameters": {"n": 1, "strength": strength},
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(GLOBAL_STYLE_API, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")
            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def query_task_result(task_id: str) -> str:
    """
    通过任务ID查询AI图像生成结果(轮询直到完成)

    参数:
        task_id: create_image函数返回的任务ID
        api_key: DashScope API访问密钥

    返回:
        任务成功时生成的图片URL

    异常:
        ValueError: 任务ID无效或任务失败
        httpx.RequestError: 网络请求失败
        TimeoutError: 超过最大重试次数
    """
    QUERY_API = "https://dashscope.aliyuncs.com/api/v1/tasks/"

    api_key = get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(f"{QUERY_API}{task_id}", headers=headers)
            response.raise_for_status()
            result = response.json()

            status = result["output"]["task_status"]
            if status == "SUCCEEDED":
                return result["output"]["results"][0]["url"]
            elif status in ["FAILED", "CANCELED"]:
                raise ValueError(f"任务正在处理中请等待，状态码: {status}")

        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def text_to_speech(text: str, voice: str = "Chelsie") -> str:
    """
    通过通义千问语音合成大模型生成语音

    Args:
        text: 要合成的文本内容
        voice: 音色选择 (Cherry/Serena/Ethan/Chelsie)，默认Chelsie

    Returns:
        生成的音频URL

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """
    TTS_API = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"

    api_key = get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {"model": "qwen-tts", "input": {"text": text, "voice": voice}}

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(TTS_API, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            if "output" not in result or "audio" not in result["output"]:
                raise ValueError("API响应格式不正确")

            return result["output"]["audio"]["url"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def text_to_video(
    prompt: str,
    model: str = "wanx2.1-t2v-plus",
    size: str = "1280*720",
    prompt_extend: bool = True,
    seed: Optional[int] = None,
) -> str:
    """
    文生视频 - 通过文本提示生成视频

    Args:
        prompt: 文本描述要生成的视频内容
        model: 模型ID，默认"wanx2.1-t2v-turbo"
        size: 视频分辨率，默认"1280*720"。支持480P和720P档位分辨率
        prompt_extend: 是否开启prompt智能改写，默认True
        seed: 随机数种子，可选

    Returns:
        task_id: 用于查询任务结果的ID

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    parameters = {
        "size": size,
        "duration": 5,
        "prompt_extend": prompt_extend,
    }
    if seed is not None:
        parameters["seed"] = seed

    data = {
        "model": model,
        "input": {"prompt": prompt},
        "parameters": parameters,
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")

            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def image_to_video(
    prompt: str,
    img_url: str,
    model: str = "wanx2.1-i2v-plus",
    resolution: str = "720P",
    prompt_extend: bool = True,
    seed: Optional[int] = None,
    template: str = "",
) -> str:
    """
    图生视频 - 基于图片生成视频

    Args:
        prompt: 文本描述要生成的视频内容
        img_url: 基础图片URL
        model: 模型ID，默认"wanx2.1-i2v-turbo"
        resolution: 视频分辨率，默认"720P"
        prompt_extend: 是否扩展提示，默认True
        seed: 随机数种子，可选
        template:指定视频特效模板的名称，用于应用特定视觉效果。若未填写，则表示不使用任何视频特效。
                当前支持的模板：
                1. flying: 魔法悬浮
                2. carousel: 时光木马
                3. squish: 解压捏捏
                4. rotation: 转圈圈
                5. poke: 戳戳乐
                6. inflate: 气球膨胀
                7. dissolve: 分子扩散
                8. singleheart: 爱你呦
                9. dance1: 摇摆时刻
                10.dance2: 头号甩舞
                11.dance3: 星摇时刻
                12.mermaid: 人鱼觉醒
                13.graduation: 学术加冕
                14.dragon: 巨兽追袭
                15.money: 财从天降
                15.rose: 赠人玫瑰
                17.crystalrose: 闪亮玫瑰
                18.hug: 爱的抱抱
                19.frenchkiss: 唇齿相依
                20.coupleheart: 双倍心动

    Returns:
        task_id: 用于查询任务结果的ID

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": model,
        "input": {"prompt": prompt, "img_url": img_url},
        "parameters": {
            "resolution": resolution,
            "prompt_extend": prompt_extend,
            "duration": 5,
        },
    }
    if seed is not None:
        data["parameters"]["seed"] = seed

    if template:
        data["input"]["template"] = template

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")

            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def image_to_video_with_frames(
    first_frame_url: str,
    last_frame_url: str,
    prompt: str = "",
    model: str = "wanx2.1-kf2v-plus",
    resolution: str = "720P",
    seed: Optional[int] = None,
) -> str:
    """
    基于首尾帧生成视频 - 通过首帧和尾帧图片生成视频

    Args:
        first_frame_url: 首帧图片URL
        last_frame_url: 尾帧图片URL
        prompt: 可选文本提示词，描述视频内容
        model: 模型ID，默认"wanx2.1-i2v-plus"
        resolution: 视频分辨率，默认"720P"
        seed: 随机数种子，可选

    Returns:
        task_id: 用于查询任务结果的ID

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2video/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": model,
        "input": {
            "first_frame_url": first_frame_url,
            "last_frame_url": last_frame_url,
            "prompt": prompt,
        },
        "parameters": {
            "resolution": resolution,
            "seed": seed,
            "duration": 5,
        },
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("API响应格式不正确")

            return result["output"]["task_id"]
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def multi_image_reference(
    prompt: str,
    ref_images_url: List[str],
    obj_or_bg: Optional[List[str]] = None,
    size: str = "1280*720",
    prompt_extend: bool = True,
    seed: Optional[int] = None,
    watermark: bool = False,
) -> str:
    """
    多图参考 - 基于多张参考图像生成视频

    Args:
        prompt: 视频内容描述文本
        ref_images_url: 参考图像URL列表(1-3张)
        obj_or_bg: 标识每张图像用途["obj"(主体)或"bg"(背景)]
        size: 视频分辨率(默认"1280*720")
        prompt_extend: 是否开启prompt智能改写(默认True)
        seed: 随机数种子(可选)
        watermark: 是否添加水印(默认False)

    Returns:
        任务ID

    Raises:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-vace-plus",
        "input": {
            "function": "image_reference",
            "prompt": prompt,
            "ref_images_url": ref_images_url[:3],  # 最多3张
        },
        "parameters": {
            "size": size,
            "prompt_extend": prompt_extend,
            "watermark": watermark,
        },
    }

    if obj_or_bg:
        data["parameters"]["obj_or_bg"] = obj_or_bg

    if seed is not None:
        data["parameters"]["seed"] = seed

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()

            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("Invalid API response format")

            return result["output"]["task_id"]

        except httpx.RequestError as e:
            raise httpx.RequestError(f"Request failed: {str(e)}")


@mcp.tool()
async def video_repainting(
    prompt: str,
    video_url: str,
    control_condition: str = "depth",
    ref_images_url: Optional[List[str]] = None,
    strength: float = 1.0,
    prompt_extend: bool = False,
    seed: Optional[int] = None,
    watermark: bool = False,
) -> str:
    """
    视频重绘 - 基于输入视频生成新风格视频

    Args:
        prompt: 视频内容描述文本
        video_url: 输入视频URL
        control_condition: 特征提取方式(默认"depth")
        ref_images_url: 参考图像URL列表(可选)
        strength: 控制强度(默认1.0)
        prompt_extend: 是否开启prompt智能改写(默认False)
        seed: 随机数种子(可选)
        watermark: 是否添加水印(默认False)

    Returns:
        任务ID

    Raises:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-vace-plus",
        "input": {
            "function": "video_repainting",
            "prompt": prompt,
            "video_url": video_url,
        },
        "parameters": {
            "control_condition": control_condition,
            "strength": strength,
            "prompt_extend": prompt_extend,
            "watermark": watermark,
        },
    }

    if ref_images_url:
        data["input"]["ref_images_url"] = ref_images_url[:1]  # 最多1张

    if seed is not None:
        data["parameters"]["seed"] = seed

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()

            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("Invalid API response format")

            return result["output"]["task_id"]

        except httpx.RequestError as e:
            raise httpx.RequestError(f"Request failed: {str(e)}")


@mcp.tool()
async def video_local_edit(
    prompt: str,
    video_url: str,
    mask_image_url: str,
    mask_frame_id: int = 1,
    ref_images_url: Optional[List[str]] = None,
    control_condition: Optional[str] = None,
    mask_type: str = "tracking",
    expand_ratio: float = 0.05,
    prompt_extend: bool = False,
    seed: Optional[int] = None,
    watermark: bool = False,
) -> str:
    """
    局部编辑 - 对视频指定区域进行编辑

    Args:
        prompt: 编辑内容描述文本
        video_url: 输入视频URL
        mask_image_url: 掩码图像URL
        mask_frame_id: 掩码目标帧ID(默认1)
        ref_images_url: 参考图像URL列表(可选)
        control_condition: 特征提取方式(可选)
        mask_type: 掩码类型(默认"tracking")
        expand_ratio: 扩展比例(默认0.05)
        prompt_extend: 是否开启prompt智能改写(默认False)
        seed: 随机数种子(可选)
        watermark: 是否添加水印(默认False)

    Returns:
        任务ID

    Raises:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-vace-plus",
        "input": {
            "function": "video_edit",
            "prompt": prompt,
            "video_url": video_url,
            "mask_image_url": mask_image_url,
            "mask_frame_id": mask_frame_id,
        },
        "parameters": {
            "mask_type": mask_type,
            "expand_ratio": expand_ratio,
            "prompt_extend": prompt_extend,
            "watermark": watermark,
        },
    }

    if ref_images_url:
        data["input"]["ref_images_url"] = ref_images_url[:1]  # 最多1张

    if control_condition:
        data["parameters"]["control_condition"] = control_condition

    if seed is not None:
        data["parameters"]["seed"] = seed

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()

            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("Invalid API response format")

            return result["output"]["task_id"]

        except httpx.RequestError as e:
            raise httpx.RequestError(f"Request failed: {str(e)}")


@mcp.tool()
async def video_extension(
    prompt: str,
    first_frame_url: Optional[str] = None,
    last_frame_url: Optional[str] = None,
    first_clip_url: Optional[str] = None,
    last_clip_url: Optional[str] = None,
    video_url: Optional[str] = None,
    control_condition: Optional[str] = None,
    prompt_extend: bool = False,
    seed: Optional[int] = None,
    watermark: bool = False,
) -> str:
    """
    视频延展 - 基于图像或视频生成延续性内容

    Args:
        prompt: 视频内容描述文本
        first_frame_url: 首帧图像URL(可选)
        last_frame_url: 尾帧图像URL(可选)
        first_clip_url: 首段视频URL(可选)
        last_clip_url: 尾段视频URL(可选)
        video_url: 参考视频URL(可选)
        control_condition: 特征提取方式(可选)
        prompt_extend: 是否开启prompt智能改写(默认False)
        seed: 随机数种子(可选)
        watermark: 是否添加水印(默认False)

    Returns:
        任务ID

    Raises:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-vace-plus",
        "input": {"function": "video_extension", "prompt": prompt},
        "parameters": {"prompt_extend": prompt_extend, "watermark": watermark},
    }

    if first_frame_url:
        data["input"]["first_frame_url"] = first_frame_url
    if last_frame_url:
        data["input"]["last_frame_url"] = last_frame_url
    if first_clip_url:
        data["input"]["first_clip_url"] = first_clip_url
    if last_clip_url:
        data["input"]["last_clip_url"] = last_clip_url
    if video_url:
        data["input"]["video_url"] = video_url
    if control_condition:
        data["parameters"]["control_condition"] = control_condition
    if seed is not None:
        data["parameters"]["seed"] = seed

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()

            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("Invalid API response format")

            return result["output"]["task_id"]

        except httpx.RequestError as e:
            raise httpx.RequestError(f"Request failed: {str(e)}")


@mcp.tool()
async def video_outpainting(
    prompt: str,
    video_url: str,
    top_scale: float = 1.0,
    bottom_scale: float = 1.0,
    left_scale: float = 1.0,
    right_scale: float = 1.0,
    prompt_extend: bool = False,
    seed: Optional[int] = None,
    watermark: bool = False,
) -> str:
    """
    视频画面扩展 - 在指定方向扩展视频画面

    Args:
        prompt: 扩展内容描述文本
        video_url: 输入视频URL
        top_scale: 顶部扩展比例(默认1.0)
        bottom_scale: 底部扩展比例(默认1.0)
        left_scale: 左侧扩展比例(默认1.0)
        right_scale: 右侧扩展比例(默认1.0)
        prompt_extend: 是否开启prompt智能改写(默认False)
        seed: 随机数种子(可选)
        watermark: 是否添加水印(默认False)

    Returns:
        任务ID

    Raises:
        ValueError: API密钥无效或响应格式错误
        httpx.RequestError: 网络请求失败
    """
    API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation/video-synthesis"

    api_key = get_api_key()
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    data = {
        "model": "wanx2.1-vace-plus",
        "input": {
            "function": "video_outpainting",
            "prompt": prompt,
            "video_url": video_url,
        },
        "parameters": {
            "top_scale": top_scale,
            "bottom_scale": bottom_scale,
            "left_scale": left_scale,
            "right_scale": right_scale,
            "prompt_extend": prompt_extend,
            "watermark": watermark,
        },
    }

    if seed is not None:
        data["parameters"]["seed"] = seed

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(API_URL, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()

            if "output" not in result or "task_id" not in result["output"]:
                raise ValueError("Invalid API response format")

            return result["output"]["task_id"]

        except httpx.RequestError as e:
            raise httpx.RequestError(f"Request failed: {str(e)}")


@mcp.tool()
async def get_video_result(task_id: str) -> str:
    """
    查询视频生成结果，轮询直到任务完成

    Args:
        task_id: 任务ID

    Returns:
        video_url: 生成的视频URL

    Raises:
        ValueError: 当任务ID无效或任务失败时
        httpx.RequestError: 网络请求失败时
        TimeoutError: 当任务超过最大重试次数时
    """
    QUERY_API = "https://dashscope.aliyuncs.com/api/v1/tasks/"

    api_key = get_api_key()
    headers = {"Authorization": f"Bearer {api_key}"}

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(f"{QUERY_API}{task_id}", headers=headers)
            response.raise_for_status()
            result = response.json()

            status = result["output"]["task_status"]
            if status == "SUCCEEDED":
                return result["output"]["video_url"]
            elif status in ["FAILED", "CANCELED"]:
                raise ValueError(f"任务正在处理中请等待，状态码: {status}")

        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
